package cn.wbu.zyh.websocket;

import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.io.IOException;
import java.net.URLDecoder;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class YjsWebSocketHandler extends BinaryWebSocketHandler {

    // 使用 Map 管理不同 noteId 的 WebSocket sessions
    private static final Map<String, Set<WebSocketSession>> noteSessions = new ConcurrentHashMap<>();

    // 每个会话的活跃状态
    private static final Map<String, Boolean> sessionActive = new ConcurrentHashMap<>();

    // 为每个会话使用单独的锁，避免并发发送消息
    private static final Map<String, Lock> sessionLocks = new ConcurrentHashMap<>();

    // 使用线程池替代为每个消息创建新线程
    private static final ExecutorService messageExecutor = Executors.newFixedThreadPool(20);

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        try {
            // 从 URI 中提取 noteId
            String path = session.getUri().getPath();
            String noteId = path.substring(path.lastIndexOf('/') + 1);

            // 从查询参数中获取用户信息
            String query = session.getUri().getQuery();
            String userName = "anonymous";
            String sessionId = session.getId();

            if (query != null && !query.isEmpty()) {
                String[] params = query.split("&");
                for (String param : params) {
                    String[] keyValue = param.split("=");
                    if (keyValue.length == 2) {
                        if (keyValue[0].equals("userName")) {
                            // URL解码用户名，处理中文和特殊字符
                            userName = URLDecoder.decode(keyValue[1], StandardCharsets.UTF_8.name());
                        } else if (keyValue[0].equals("sessionId")) {
                            sessionId = keyValue[1];
                        }
                    }
                }
            }

            // 存储会话信息
            session.getAttributes().put("noteId", noteId);
            session.getAttributes().put("userName", userName);
            session.getAttributes().put("sessionId", sessionId);

            // 为会话创建锁
            sessionLocks.put(session.getId(), new ReentrantLock());

            // 标记会话为活跃
            sessionActive.put(session.getId(), true);

            // 将会话添加到noteId对应集合
            Set<WebSocketSession> sessions = noteSessions.computeIfAbsent(noteId, k -> new CopyOnWriteArraySet<>());
            sessions.add(session);

            System.out.println("Connected to noteId: " + noteId + ", user: " + userName + ", session: " + sessionId);

            // 发送连接成功消息给客户端
            sendConnectionStatusMessage(session, true);
        } catch (Exception e) {
            System.err.println("Error in afterConnectionEstablished: " + e.getMessage());
            e.printStackTrace();
        }
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        try {
            // 从会话中获取 noteId
            String noteId = (String) session.getAttributes().getOrDefault("noteId", "");

            if (noteId.isEmpty()) {
                System.err.println("Session has no noteId: " + session.getId());
                return;
            }

            // 读取消息内容
            ByteBuffer payload = message.getPayload();
            byte[] bytes = new byte[payload.remaining()];
            payload.get(bytes);

            // 获取特定 noteId 的会话集合
            Set<WebSocketSession> sessions = noteSessions.getOrDefault(noteId, new CopyOnWriteArraySet<>());

            // 广播消息到同一笔记的其他会话
            for (WebSocketSession recipient : sessions) {
                if (!recipient.isOpen() || recipient.getId().equals(session.getId()) ||
                        !Boolean.TRUE.equals(sessionActive.get(recipient.getId()))) {
                    continue;
                }

                // 获取此会话的锁
                Lock lock = sessionLocks.get(recipient.getId());
                if (lock == null) {
                    lock = new ReentrantLock();
                    sessionLocks.put(recipient.getId(), lock);
                }

                // 使用线程池异步发送消息
                final Lock finalLock = lock;
                messageExecutor.execute(() -> sendMessageWithLock(recipient, bytes, finalLock));
            }
        } catch (Exception e) {
            System.err.println("Error in handleBinaryMessage: " + e.getMessage());
            e.printStackTrace();
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        String sessionId = session.getId();
        String noteId = (String) session.getAttributes().getOrDefault("noteId", "");
        String userName = (String) session.getAttributes().getOrDefault("userName", "anonymous");

        System.err.println("Transport error for user: " + userName + ", session: " + sessionId + ", error: " + exception.getMessage());

        // 标记会话为非活跃
        sessionActive.put(sessionId, false);

        // 尝试关闭会话
        try {
            if (session.isOpen()) {
                session.close(CloseStatus.SERVER_ERROR);
            }
        } catch (IOException e) {
            System.err.println("Error closing session after transport error: " + e.getMessage());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        try {
            // 从会话中获取信息
            String noteId = (String) session.getAttributes().getOrDefault("noteId", "");
            String userName = (String) session.getAttributes().getOrDefault("userName", "anonymous");
            String sessionId = (String) session.getAttributes().getOrDefault("sessionId", session.getId());

            // 移除会话
            if (!noteId.isEmpty()) {
                Set<WebSocketSession> sessions = noteSessions.get(noteId);
                if (sessions != null) {
                    sessions.remove(session);

                    // 通知其他用户此用户已离开
                    notifyUserDisconnected(sessions, userName, sessionId);

                    if (sessions.isEmpty()) {
                        noteSessions.remove(noteId);
                    }
                }
            }

            // 移除会话锁和活跃状态
            sessionLocks.remove(session.getId());
            sessionActive.remove(session.getId());

            System.out.println("Disconnected from noteId " + noteId + ", user: " + userName + ", session: " + sessionId);
        } catch (Exception e) {
            System.err.println("Error in afterConnectionClosed: " + e.getMessage());
            e.printStackTrace();
        }
    }

    // 使用锁发送消息，确保同一会话的消息发送是线性的
    private void sendMessageWithLock(WebSocketSession session, byte[] data, Lock lock) {
        lock.lock();
        try {
            if (session.isOpen() && Boolean.TRUE.equals(sessionActive.get(session.getId()))) {
                BinaryMessage binaryMessage = new BinaryMessage(data);
                session.sendMessage(binaryMessage);
            }
        } catch (Exception e) {
            System.err.println("Error sending message to session " + session.getId() + ": " + e.getMessage());
            // 标记会话为非活跃
            sessionActive.put(session.getId(), false);
        } finally {
            lock.unlock();
        }
    }

    // 发送连接状态消息给客户端
    private void sendConnectionStatusMessage(WebSocketSession session, boolean connected) {
        try {
            String message = "{\"type\":\"connection_status\",\"connected\":" + connected + "}";
            session.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            System.err.println("Error sending connection status message: " + e.getMessage());
        }
    }

    // 通知其他用户某用户已断开连接
    private void notifyUserDisconnected(Set<WebSocketSession> sessions, String userName, String sessionId) {
        if (sessions == null || sessions.isEmpty()) {
            return;
        }

        String message = "{\"type\":\"user_disconnected\",\"userName\":\"" +
                userName.replace("\"", "\\\"") + "\",\"sessionId\":\"" + sessionId + "\"}";

        for (WebSocketSession recipient : sessions) {
            if (recipient.isOpen() && Boolean.TRUE.equals(sessionActive.get(recipient.getId()))) {
                try {
                    recipient.sendMessage(new TextMessage(message));
                } catch (IOException e) {
                    System.err.println("Error sending user disconnection notification: " + e.getMessage());
                }
            }
        }
    }
}